#--------------------------------------------#
#   该部分代码只用于看网络结构，并非测试代码
#   map测试请看get_dr_txt.py、get_gt_txt.py
#   和get_map.py
#--------------------------------------------#
import torch
import glob

from nets.efficientdet import EfficientDetBackbone
from nets.efficientnet import EfficientNet

def torchviz_demo():
    from torchviz import make_dot
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
    model = EfficientDetBackbone(80,0).to(device)
    x = torch.randn(1, 3, 768, 768)
    y = model(x)
    vise=make_dot(y, params=dict(model.named_parameters()))
    vise.view() #结果（以PDF格式保存在源程序所在文件夹）

def torchsummary_demo():
    from torchsummary import summary
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
    model = EfficientDetBackbone(80,0).to(device) 
    summary(model,input_size=(3, 768, 768))

def Thop_demo():
    from thop import profile
    input = torch.randn(1, 3, 768, 768)	
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
    model = EfficientDetBackbone(80,0).to(device)
    flops, params = profile(model, inputs=(input, ))

if __name__ == '__main__':
    #inputs = torch.randn(4, 3, 512, 512)
    #model = EfficientDetBackbone(80,0)
    #print('# generator parameters:', sum(param.numel() for param in model.parameters()))
    torchviz_demo()
